#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>

#define BUF_SIZE 1024
#define SERVER_PORT 8887

int main() {
    char buf[BUF_SIZE];
    int sock, send_sock;
    struct sockaddr_in sock_addr, send_addr;
    char other_ip[15];
    char message[BUF_SIZE];

    printf("是否为服务器 [0/1]: ");
    memset(buf, 0, sizeof(buf));
    fgets(buf, BUF_SIZE, stdin);
    buf[strlen(buf) - 1] = '\0';
    int is_server = atoi(buf);

    if (!is_server) {
        // 输入服务器IP地址
        printf("请输入对方的IP地址: ");
        memset(buf, 0, sizeof(buf));
        fgets(buf, BUF_SIZE, stdin);
        buf[strlen(buf) - 1] = '\0';
        strcpy(other_ip, buf);
    }

    // 创建套接字
    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == -1) {
        printf("创建套接字失败\n");
        return 1;
    }

    // 设置地址
    sock_addr.sin_family = AF_INET;
    sock_addr.sin_addr.s_addr = is_server ? INADDR_ANY : inet_addr(other_ip);
    sock_addr.sin_port = htons(SERVER_PORT);

    if (is_server) {
        // 绑定套接字
        if (bind(sock, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
            printf("绑定失败\n");
            return 1;
        }
        // 监听连接
        listen(sock, 1);
        printf("等待客户端连接...\n");
        // 接受连接
        int c = sizeof(struct sockaddr_in);
        send_sock = accept(sock, (struct sockaddr *)&send_addr, (socklen_t *)&c);
        if (send_sock < 0) {
            printf("连接失败\n");
            return 1;
        }
    } else {
        // 连接到服务器
        if (connect(sock, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
            printf("连接失败\n");
            return 1;
        }
        send_sock = sock;
        send_addr = sock_addr;
    }

    printf("连接成功! 现在可以开始聊天了\n");

    // 发送和接收消息
    while (1) {
        printf("发送消息: ");
        memset(buf, 0, sizeof(buf));
        fgets(buf, BUF_SIZE, stdin);
        buf[strlen(buf) - 1] = '\0';
        strcpy(message, buf);

        if (strlen(message) == 0) break;

        // 发送消息
        if (send(send_sock, message, strlen(message), 0) < 0) {
            printf("发送失败\n");
            break;
        }

        // 清空缓冲区
        memset(message, 0, sizeof(message));

        // 接收消息
        if (recv(send_sock, message, BUF_SIZE, 0) < 0) {
            printf("接收失败\n");
            break;
        }

        printf("对方: %s\n", message);
    }

    // 关闭套接字
    close(sock);
    close(send_sock);

    return 0;
}